# coding = utf-8
# @Time    : 2024-05-21  13:38:27
# @Author  : zhaosheng@nuaa.edu.cn
# @Describe: Dguard Test.

import os
import sys
import torch
import random
import argparse
import torchaudio
import numpy as np

from tqdm import tqdm
from kaldiio import WriteHelper
from wespeaker.utils.utils import set_seed
from dguard.utils.utils import get_logger
from dguard.utils.fileio import load_wav_scp
from dguard import DguardModel as dm
parser = argparse.ArgumentParser(description='Extract embeddings for evaluation.')
parser.add_argument('--exp_dir', default='', type=str, help='Exp dir')
parser.add_argument('--data', default='', type=str, help='Data dir')
parser.add_argument('--model_name', default='', type=str, help='model_name')
parser.add_argument('--key', default='', type=str, help='key')
parser.add_argument('--seed', default=123, type=int, help='seed')
parser.add_argument('--length', default=3.0, type=float, help='wav length')
parser.add_argument('--start_from', default=0.0, type=float, help='start from(s)')
parser.add_argument('--use_gpu', action='store_true', help='Use gpu or not')
parser.add_argument('--mean_embedding', action='store_true', help='Use mean embedding or not')
parser.add_argument('--gpu', nargs='+', help='GPU id to use.')



def main():
    args = parser.parse_args(sys.argv[1:])
    set_seed(args.seed)
    rank = int(os.environ['LOCAL_RANK'])
    world_size = int(os.environ['WORLD_SIZE'])
    embedding_dir = os.path.join(args.exp_dir, 'embeddings')
    os.makedirs(embedding_dir, exist_ok=True)
    logger = get_logger(fpath=os.path.join(embedding_dir, 'extract_emb.log'))
    print(f"Mean embedding: {args.mean_embedding}")
    if args.use_gpu:
        if torch.cuda.is_available():
            gpu = int(args.gpu[rank % len(args.gpu)])
            device = torch.device('cuda', gpu)
            device = f'cuda:{gpu}'
        else:
            msg = 'No cuda device is detected. Using the cpu device.'
            if rank == 0:
                logger.warning(msg)
            device = torch.device('cpu')
            device = 'cpu'
    else:
        device = torch.device('cpu')
        device = 'cpu'

    if "_and_" in args.model_name:
        MODELS = args.model_name.split("_and_")
        data = load_wav_scp(args.data)
        data_k = list(data.keys())
        local_k = data_k[rank::world_size]
        if len(local_k) == 0:
            logger.error("The number of threads exceeds the number of files")
            sys.exit()
            
        for _index in range(len(MODELS)):
            now_model = MODELS[_index]
            dguard_model = dm(embedding_model_names=[now_model],
                                device=device,
                                length=args.length,
                                start_time=args.start_from,
                                mean=args.mean_embedding,
                                apply_vad=False,
                                seed=args.seed,
                                key=args.key)
            # dguard_model.verify(args.key)
            emb_ark = os.path.join(embedding_dir, f"fusion_{now_model}_"+'xvector_%02d.ark'%rank)
            emb_scp = os.path.join(embedding_dir, f"fusion_{now_model}_"+'xvector_%02d.scp'%rank)
            if rank == 0:
                logger.info('Start extracting embeddings.')
            with WriteHelper(f'ark,scp:{emb_ark},{emb_scp}') as writer:
                for k in tqdm(local_k):
                    wav_path = data[k]
                    wav, fs = torchaudio.load(wav_path)
                    emb = dguard_model.encode(wav_file=wav_path)
                    emb = emb.reshape(-1)
                    print(f"emb shape: {emb.shape}")
                    emb = emb.detach().cpu().numpy()
                    writer(k, emb)

    else:  
        dguard_model = dm(embedding_model_names=[args.model_name],
                                device=device,
                                length=args.length,
                                start_time=args.start_from,
                                mean=args.mean_embedding,
                                apply_vad=False,
                                seed=args.seed,
                                key=args.key)
        dguard_model.verify(args.key)
        data = load_wav_scp(args.data)
        data_k = list(data.keys())
        local_k = data_k[rank::world_size]
        if len(local_k) == 0:
            msg = "The number of threads exceeds the number of files"
            logger.info(msg)
            sys.exit()

        emb_ark = os.path.join(embedding_dir, 'xvector_%02d.ark'%rank)
        emb_scp = os.path.join(embedding_dir, 'xvector_%02d.scp'%rank)
        if rank == 0:
            logger.info('Start extracting embeddings.')
    
        with WriteHelper(f'ark,scp:{emb_ark},{emb_scp}') as writer:
            for k in tqdm(local_k):
                wav_path = data[k]
                embs = dguard_model.encode(wav_file=wav_path)
                embs = embs.reshape(-1)
                emb = embs.detach().cpu().numpy()
                emb = np.concatenate((emb,emb,emb),axis=0)
                emb = emb.reshape(-1)
                writer(k, emb)

if __name__ == "__main__":
    main()
